functions {
   real pow_sum (vector x, real z){
    real y = 0;
    
    for (i in 1:num_elements(x)){
      y = y + x[i]^z;
    }
    return y;
   }
   
  real sqsum (vector x){
    real y = 0;
    
    for (i in 1:num_elements(x)){
      y = y + x[i]*x[i];
    }
    return y;
   } 
   
   real calc_cor (vector x, vector y){
     real r;
     r = sum((x - mean(x)).*(y - mean(y))) / sqrt(sqsum(x - mean(x))*sqsum(y - mean(y)));
     return r;
   }
}
data {
  int<lower=1> I;               // proposals
  int<lower=1> J;               // legislators
  int<lower=1> D;               // number of latent dimensions
  int<lower=1> K;               // number of exp. variables
  int<lower=0,upper=1> Y[J,I];  // binary y
  real theta_in[J,K,D];         // explanatory variables for ideal points
  matrix[J,D] theta_true;
  int<lower=0,upper=1> incl_nonsep;  // indicator to use separable or nonseparable model
  int<lower=0,upper=1> incl_logl;  // indicator to calculate log likelihood
}
parameters {
  ordered[D] sal_free; // salience of dimensions - free
  real<lower=0> tau; // scale of weights matrix overall
  real<lower=0> tau_int; // scale of weights matrix internally
  matrix[J,D] theta_free; // free ideal points
  matrix[I,D] bk_free; // free discrimination
  vector[I] ak; // difficulty
  cholesky_factor_corr[D] L; // correlation component of weight matrix
  row_vector<lower=0>[K] gamma[D]; // coefficients for linear regression on ideal points, lower bound to prevent switching
}
transformed parameters {
  corr_matrix[D] DL; // correlation component of weight matrix
  cov_matrix[D] weights; // final weight matrix
  matrix[J,D] theta_t;
  matrix[I,D] bk; // standardized discrimination (var 1)
  matrix[J,D] theta; // standardized ideal points (mean 0, var 1)
  simplex[D] sal = softmax(sort_desc(sal_free)); // salience of dimensions - rescaled to simplex
  DL = L * L';
  
  {
    vector[D] sd_theta;
    vector[D] sd_bk;
    vector[D] mean_theta;
    
    // needs this constraint of sum(abs(...)) == 1 here so that weights is identified internally.
    // otherwise, same variance can be gotten by varying either var or cov within the final matrix
    if (incl_nonsep == 1){
      weights = quad_form_diag(DL, sqrt(sal * tau_int)) / sum(fabs(to_vector(quad_form_diag(DL, sqrt(sal * tau_int))))) * tau^2;
    } else {
      weights = quad_form_diag(diag_matrix(rep_vector(1.00, D)), sqrt(sal * tau));
    }
    for (d in 1:D){
      theta_t[,d] = to_matrix(theta_in[,,d]) * to_vector(gamma[d]) + theta_free[,d];
    }
    
    for (d in 1:D){
      sd_theta[d] = sd(theta_t[,d]);
      sd_bk[d] = sd(bk_free[,d]);
      mean_theta[d] = mean(theta_t[,d]);
    }

    for (d in 1:D){
      theta[,d] = (theta_t[,d] - mean_theta[d]) / sd_theta[d]; 
      bk[,d] = (bk_free[,d]) / sd_bk[d];
    }
  }
}
model {
  to_vector(bk_free) ~ std_normal();
  ak ~ std_normal();
  sal_free ~ std_normal();
  tau ~ std_normal();
  tau_int ~ normal(D*1.00, 1); // in expectation, the 'internal' weights part has a variance of 2 along the diagonals, before being rescaled by tau
  to_vector(theta_free) ~ std_normal();
  for (d in 1:D){
    gamma[d] ~ normal(1, 1);
  }
  L ~ lkj_corr_cholesky(1); // uniform dist
  
  {
    matrix[J,I] ystar;
    ystar = (theta * weights * bk');
    for (i in 1:I){
      Y[,i] ~ bernoulli(Phi_approx(ystar[,i] + ak[i]));
    }
  }
  
}
generated quantities{
  vector[D] r_sq; // share of var in final theta explained by lpm
  vector[D] r_true; // correlation of final theta with true theta
  real r_dim; // correlation of final theta with true theta
  real m_prop_compl; // mean proposal complexity
  real sd_prop_compl; // sd proposal complexity
  real p_share; // share of correct votes
  matrix[J,D] theta_pred; // predicted values based on lpm
  matrix[J,D] theta_resid; // residual values
  row_vector[K] gamma_post[D]; // coef from lpm on rescaled theta
  vector[I*J*incl_logl] log_lik;
  
  r_dim = calc_cor(theta[,1], theta[,2]);
  
  for (d in 1:D){
    theta_pred[,d] = to_matrix(theta_in[,,d]) * to_vector(gamma[d]); 
    theta_resid[,d] = (theta_pred[,d] - mean(theta_pred[,d])) - (theta_t[,d] - mean(theta_t[,d]));
    r_sq[d] = sd(theta_pred[,d]) * sd(theta_pred[,d]) / (sd(theta_pred[,d]) * sd(theta_pred[,d]) + sd(theta_resid[,d]) * sd(theta_resid[,d]));
    r_true[d] = calc_cor(theta_true[,d], theta[,d]);
    gamma_post[d] = to_row_vector((to_matrix(theta_in[,,d])' * to_matrix(theta_in[,,d])) \ to_matrix(theta_in[,,d])' * theta[,d]);
  }
  {
    vector[I] prop_compl;
    for (i in 1:I) prop_compl[i] = (pow_sum(bk[i,]', 2.00))^2 / pow_sum(bk[i,]', 4.00); // Hofmann, R. J. (1977). Indices descriptive of factor complexity. Journal of General Psychology, 96(1), 103â€“110. 
    m_prop_compl = mean(prop_compl);
    sd_prop_compl = sd(prop_compl);
  }
  if (incl_logl == 1){
    {
      matrix[J,I] ystar;
      real pred = 0.0;
      for (i in 1:I){
        for (j in 1:J){
         ystar[j,i] = bernoulli_lpmf(Y[j,i] | Phi_approx((theta[j,] * weights * bk[i,]') + ak[i]));
         if (Y[j,i] == bernoulli_rng(Phi_approx((theta[j,] * weights * bk[i,]') + ak[i]))){
           pred += 1.00;
         }
        }
      }
      log_lik = to_vector(ystar);
      p_share = pred / (I*J*1.00);
    }
  }
}